module windump.plugins.pstore.pstore;

import std.c.windows.com;
import std.utf;
import std.string;

import windump.utf;
import windump.w32api;
import windump.registry;
import windump.dll;

import windump.scanning.detector;
import windump.plugins.ie.iecommon;

import windump.plugins.pstore.pstoreclib;


alias DWORD PST_PROVIDERCAPABILITIES; 
alias DWORD PST_REPOSITORYCAPABILITIES;
alias DWORD PST_KEY;
alias DWORD PST_ACCESSMODE;
alias DWORD PST_ACCESSCLAUSETYPE;
alias GUID UUID;

enum : uint 
{
    PST_KEY_CURRENT_USER = 0x00000000,
    PST_KEY_LOCAL_MACHINE = 0x00000001
}



class Pstore
{
    private:
        PST_KEY		m_key;
        IPStore 	m_ps;

        DllImport!("pstorec.dll", "PStoreCreateInstance", 
                HRESULT function( IPStore* ppProvider,
                    PST_PROVIDERID* pProviderID,
                    void* pReserved,
                    DWORD dwFlags)
                ) m_fpPStoreCreateInstance;


    public:	
        class TypeEnumerator
        {		
            private IEnumPStoreTypes	m_types;

            public this(Pstore ps)                 
            {
                ps.getIPStore().EnumTypes(ps.getKey(), 0, &m_types);
            }

            public ~this()
            {
                m_types.Release();
            }

            public bool Next(out GUID id)                
            {
                DWORD celt = 1;
                return m_types.Next(celt, id, celt) == S_OK;
            }

        };

        class SubtypeEnumerator
        {		
            private IEnumPStoreTypes	m_types;

            public this(Pstore ps, GUID typeID)
            {
                ps.getIPStore().EnumSubtypes(ps.getKey(), typeID, 0, &m_types);
            }

            public    ~this()
            {
                m_types.Release();
            }

            public bool Next(out GUID id)                
            {
                DWORD celt = 1;
                return m_types.Next(celt, id, celt) == S_OK;
            }

        };

        class ItemEnumerator
        {		
            private IEnumPStoreItems	m_items;

            public this(Pstore ps, GUID typeID, GUID subtypeID)
            {
                ps.getIPStore().EnumItems(ps.getKey(), typeID, subtypeID, 0, &m_items);
            }

            public ~this()
            {
                m_items.Release();
            }

            public bool Next(out wstring itemName)
            {
                wchar* wname = null;			
                DWORD celt = 1;
                bool ret = m_items.Next(celt, wname, celt) == S_OK;
                if(!ret)return ret;
                // I don't like GC
//                itemName.length = strlen(wname);
                itemName = fromUtf16z(wname);

                scope(exit)
                    CoTaskMemFree(wname);

                return ret;
            }

        };


    public:
        this(PST_KEY key = PST_KEY_CURRENT_USER)
        {
            HRESULT hr = CoInitialize(null);		//Initialize OLE
            if (FAILED(hr))
                throw new Exception("Failed to Initialize OLE");
                
            hr = m_fpPStoreCreateInstance(&m_ps, null, null, 0);
            if(FAILED(hr))
                throw new Exception("Failed to load pstorec.dll");

            scope(failure)  {
                CoFreeUnusedLibraries();
                CoUninitialize();
                m_ps = null;            
            }
        }

        ~this()
        {            
            close();
        }

        public void close()
        {
            if(m_ps !is null) {
                CoFreeUnusedLibraries();
                CoUninitialize();
                m_ps = null;            
            }
        }

        ubyte[] readItem(GUID typeID, GUID subtypeID, string itemName)
        {
            return readItem(typeID, subtypeID, toUTF16z(itemName));
        }


        ubyte[] readItem(GUID typeID, GUID subtypeID, wchar* itemName)
        {
            uint dataSize = 0;
            ubyte* bdata = null;
            _PST_PROMPTINFO *pi = null;
            HRESULT hr = m_ps.ReadItem(m_key, typeID, subtypeID, itemName, 
                    dataSize, bdata, pi, PST_PROMPT_QUERY);

            if(hr != S_OK)
                throw new Exception("Failed to read Item");

            ubyte[] data;
            data.length = dataSize;
            data[] = bdata[0..dataSize];

            scope(success)
                CoTaskMemFree(bdata);

            return data;
        }

        PST_KEY getKey() { return m_key; }

        IPStore getIPStore() { assert(!(m_ps is null)); return m_ps; }
        
        int enumerate(int delegate(GUID /* type ID*/, GUID /* subtype ID*/, 
                    string /* itemName */) proc)
        { 
            assert(!(proc == null));

            GUID subtype;
            GUID type;
            int result = 0;


            auto typeEnum = new TypeEnumerator(this);
            while(typeEnum.Next(type))
            {	                
                auto subtypeEnum = new Pstore.SubtypeEnumerator(this, type);
                while(subtypeEnum.Next(subtype))
                {
                    auto itemEnum = new Pstore.ItemEnumerator(this, type, subtype);		
                    wstring itemName;
                    while(itemEnum.Next(itemName))
                    {					
                        result = proc(type, subtype, toUTF8(itemName));
                        if(result)break;
                    }
                    if(result)break;

                }
                if(result)break;
            }
            return result;
        }

};


class PStoreDetector : Detector!(PStoreDetector, "system/pstore", "table")
{
    private enum ResourceType : uint
    {
        OutlookAccount = 0x220D5CC1,
        IEProtectedSite = 0x5E7E8100,
        IEAutoComplete = 0xE161255A,
        OutlookExpressID =  0x89C39569,
        MsnSignup = 0xB9819C52
    }

    private const string RegexpIE6ProtectedSite = 
        `(+*):(+*)\x00`;        

    private struct KnownItems
    {
        string[string] outlookPasswords;
    }

    private struct OutlookAccount
    {
        string popHost;
        string popAccount;
        string popPassword;
    }

    private struct IEAutoCompletePair
    {
        string user;
        string password;
    }

    private IEAutoCompletePair[] decodeIEAutoComplete(ubyte[] stringIndex, ubyte[] stringData) 
    {
        IEAutoCompletePair[] ret;
        StringIndexHeader*  hdr = 
            cast(StringIndexHeader*)stringIndex.ptr;
        StringIndexEntry*   entries = 
            cast(StringIndexEntry*)(stringIndex.ptr + hdr.structureSize);
        
        for(uint i = 0; i < hdr.entriesCount; i += 2)
        {
            IEAutoCompletePair pair;
            wchar* ptr = null;
            ptr = cast(wchar*)(stringData.ptr + entries[i].dataOffset);
            pair.user = toUTF8( ptr[0 .. entries[i].dataSize]);
            ptr = cast(wchar*)(stringData.ptr + entries[i + 1].dataOffset);
            pair.password = toUTF8( ptr[0 .. entries[i + 1].dataSize]);
            ret ~= pair;
        }

        return ret;
    }

    public this()
    {
    }

    private OutlookAccount[] getOutlookAccounts()
    {
        Key hkcu = Registry.currentUser;
        scope Key k = hkcu.openSubkey(`Software\Microsoft\Internet Account Manager\Accounts`);

        OutlookAccount[] ret;
        for(uint x = 0; x < k.subkeyCount; x++)
        {
            scope subkey = k.openSubkey(x);
            OutlookAccount oa;
            try
            {
                oa.popHost = subkey.getValue("POP3 Server").value_SZ;

                // FIXME A big big big BUG within std.windows.registry !!! 
                // This is a working around for std.registry.Key.value*'s bug
                oa.popAccount = subkey.getValue("POP3 User Name").value_SZ;

                // ptr + 2 means skip the length field
                oa.popPassword = 
                    windump.utf.toUtf8(cast(wchar*)(subkey.getValue("POP3 Password2").value_BINARY.ptr + 2)).dup;
            }
            catch(RegistryException e) 
            { 
                continue;
            }

            ret ~= oa;
        }

        return ret;
    }

    public override ResultItem[] update()
    {
        auto ps = new Pstore();
        scope(exit) ps.close();

        OutlookAccount[] accounts;
        try
        {
            accounts = getOutlookAccounts();
        }
        catch(Exception e)
        {
            accounts.length = 0;
        }

        //KnownItems ki;
        ResultItem[] ret;

        int proc(GUID typeID, GUID subtypeID, string name)
        {
            switch(typeID.Data1)
            {
                case ResourceType.OutlookAccount:
                    ubyte[] data;
                    data = ps.readItem(typeID, subtypeID, name);
                    string pp = std.string.toString(cast(char*)data.ptr);
                    //find matched account
                    foreach(OutlookAccount acc; accounts)
                    {
                        if(acc.popPassword == name){
                            ResultItem item;
                            item.provider = this.name();
                            item.host = acc.popHost; 
                            item.account = acc.popAccount;
                            item.password = pp;
                            ret ~= item;
                        }
                    }
                    break;

                case ResourceType.IEProtectedSite:
					//FIXME:暂时不支持 FTP 帐号

					if(cmp(name[0..6], "DPAPI:") == 0)
						break;

                    ubyte[] data;
                    data = ps.readItem(typeID, subtypeID, name);

                    char* str = cast(char*)data.ptr;                    

                    int pos = find(str[0 .. data.length - 1], ':'); 
                    if(pos <= 0 && data.length < 3) break;

                    ResultItem item;
                    item.provider = this.name();

                    item.host = name;
                    item.account = str[0 .. pos];
                    item.password = str[pos + 1 .. data.length - 1];
                    ret ~= item;
                    break;

                case ResourceType.IEAutoComplete:
                    int pos = name.find(`:StringIndex`);
                    if(pos < 0) break;
                    string res = name[0..pos];
                    ubyte[] stringIndex = 
                        ps.readItem(typeID, subtypeID, name);
                    ubyte[] stringData = 
                        ps.readItem(typeID, subtypeID, res ~ `:StringData`);
                    IEAutoCompletePair[] pairs = 
                        decodeIEAutoComplete(stringIndex, stringData);
                    foreach(IEAutoCompletePair pair; pairs)
                    {
                        ResultItem item;
                        item.provider = this.name();
                        item.host = res;
                        item.account = pair.user;
                        item.password = pair.password;
                        ret ~= item;
                    }

                    break;

                    /*
                case ResourceType.OutlookExpressID:
                    //resType = "OutlookExpressID";
                    break;
                    */

                default:
                    //unknown type
            }

            return 0;
        }

        ps.enumerate(&proc);

        return ret;
    }

}
